from transformers import AutoTokenizer, AutoModelForCausalLM
import torch



# 本地模型的路径
model_path = r"D:\deepSeek1.5B"

# 加载分词器
tokenizer = AutoTokenizer.from_pretrained(model_path)

# 加载模型
device = "cuda" if torch.cuda.is_available() else "cpu"
model = AutoModelForCausalLM.from_pretrained(model_path).to(device)

while True:
    input_text = input("请输入待扩写的文本:")

    input_ids = tokenizer.encode(input_text, return_tensors="pt").to(device)
    output = model.generate(input_ids, max_length=1000, num_return_sequences=1)
    generated_text = tokenizer.decode(output[0], skip_special_tokens=True)
    print(generated_text)

# @app.route('/', methods=['GET', 'POST'])
# def index():
#     if request.method == 'POST':
#         input_text = request.form.get('input_text')
        # input_ids = tokenizer.encode(input_text, return_tensors="pt").to(device)
        # output = model.generate(input_ids, max_length=200, num_return_sequences=1)
        # generated_text = tokenizer.decode(output[0], skip_special_tokens=True)

        # 对输入文本进行分词处理
    #     inputs = tokenizer(input_text, return_tensors="pt", padding=True, truncation=True)
    #     # 手动创建注意力掩码
    #     input_ids = inputs["input_ids"]
    #     attention_mask = (input_ids != tokenizer.pad_token_id).long()
    #     # 将注意力掩码添加到输入中
    #     inputs["attention_mask"] = attention_mask
    #     # 进行模型推理
    #     with torch.no_grad():
    #         outputs = model(**inputs)
    #     # 获取预测结果
    #     generated_text = outputs.logits
    #     print(generated_text)
    #     return render_template_string('''
    #         <!DOCTYPE html>
    #         <html lang="en">
    #         <head>
    #             <meta charset="UTF-8">
    #             <title>DeepSeek交互网页</title>
    #         </head>
    #         <body>
    #             <h1>DeepSeek小说扩写</h1>
    #             <form method="post">
    #                 <label for="input_text">输入待扩写的文本:</label><br>
    #                 <textarea id="input_text" name="input_text" rows="4" cols="50">{{ input_text }}</textarea><br>
    #                 <input type="submit" value="扩写">
    #             </form>
    #             <h2>扩写结果:</h2>
    #             <p>{{ generated_text }}</p>
    #         </body>
    #         </html>
    #     ''', input_text=input_text, generated_text=generated_text)
    # return render_template_string('''
    #     <!DOCTYPE html>
    #     <html lang="en">
    #     <head>
    #         <meta charset="UTF-8">
    #         <title>DeepSeek交互网页</title>
    #     </head>
    #     <body>
    #         <h1>DeepSeek小说扩写</h1>
    #         <form method="post">
    #             <label for="input_text">输入待扩写的文本:</label><br>
    #             <textarea id="input_text" name="input_text" rows="4" cols="50"></textarea><br>
    #             <input type="submit" value="扩写">
    #         </form>
    #     </body>
    #     </html>
    # ''')


# if __name__ == '__main__':
    # app.run(debug=True)
    